//
//  PinterestLayout.swift
//  CustomizeCollectionView
//
//  Created by GZOffice_hao on 2019/4/25.
//  Copyright © 2019 VT. All rights reserved.
//

import UIKit

protocol PinterestLayoutDelegate {
//    func collectionView(_ collectionView: UICollectionView, heightForItemAt indexPath: IndexPath) -> CGFloat
    func collectionView(_ collectionView: UICollectionView, heightForPhotoAt indexPath: IndexPath, with width: CGFloat) -> CGFloat
    func collectionView(_ collectionView: UICollectionView, heightForAnnotationAt indexPath: IndexPath, with width: CGFloat) -> CGFloat
}

class PinterestLayoutAttributes: UICollectionViewLayoutAttributes {
    
    var photoHeight: CGFloat = 0
    
    override func copy(with zone: NSZone? = nil) -> Any {
        let copy = super.copy(with: zone) as! PinterestLayoutAttributes
        copy.photoHeight = photoHeight
        return copy
    }
    
    override func isEqual(_ object: Any?) -> Bool {
        if let attributes = object as? PinterestLayoutAttributes {
            if attributes.photoHeight == photoHeight {
                return super.isEqual(object)
            }
        }
        return false
    }
}

class PinterestLayout: UICollectionViewLayout {
    
    var cellPadding: CGFloat = 0
    var numberOfColumns = 1
    var delegate: PinterestLayoutDelegate!
    
    fileprivate var contentHeight: CGFloat = 0
    fileprivate var caches = [PinterestLayoutAttributes]()
    fileprivate var width: CGFloat {
        get {
            guard let collectionView = collectionView else { return 0 }
            let insets = collectionView.contentInset
            return collectionView.bounds.width - (insets.left + insets.right)
        }
    }
    
    override class var layoutAttributesClass: AnyClass {
        return PinterestLayoutAttributes.self
    }
    
    override var collectionViewContentSize: CGSize {
        return .init(width: width, height: contentHeight)
    }
    
    override func prepare() {
        guard let collectionView = collectionView else { return }
        if caches.isEmpty {
            let columnWidth = width / CGFloat(numberOfColumns)
            
            var xOffsets = [CGFloat]()
            for columns in 0..<numberOfColumns {
                xOffsets.append(CGFloat(columns) * columnWidth)
            }
            
            var yOffsets = [CGFloat](repeating: 0, count: numberOfColumns)
            var column = 0
            for item in 0..<collectionView.numberOfItems(inSection: 0) {
                let indexPath = IndexPath(item: item, section: 0)
                let width = columnWidth - 2 * cellPadding
                let photoHeight = delegate.collectionView(collectionView, heightForPhotoAt: indexPath, with: width)
                let annotationHeight = delegate.collectionView(collectionView, heightForAnnotationAt: indexPath, with: width)
                let height = cellPadding + photoHeight + annotationHeight + cellPadding
                let frame = CGRect(x: xOffsets[column], y: yOffsets[column], width: columnWidth, height: height)
                let insetFrame = frame.insetBy(dx: cellPadding, dy: cellPadding)
                let attributes = PinterestLayoutAttributes(forCellWith: indexPath)
                attributes.frame = insetFrame
                attributes.photoHeight = photoHeight
                caches.append(attributes)
                contentHeight = max(contentHeight, insetFrame.maxY)
                yOffsets[column] = yOffsets[column] + height
                column = column >= (numberOfColumns - 1) ? 0 : column + 1
            }
        }
    }
    
    override func layoutAttributesForElements(in rect: CGRect) -> [UICollectionViewLayoutAttributes]? {
        var layoutAttributes = [UICollectionViewLayoutAttributes]()
        for attributes in caches {
            if rect.intersects(attributes.frame) {
                layoutAttributes.append(attributes)
            }
        }
        return layoutAttributes
    }
}
